# -*- coding: utf-8 -*-
# @Author: bao
# @Date:   2017-03-31 15:10:31
# @Last Modified by:   bao
# @Last Modified time: 2017-05-01 14:39:52
import sys
sys.path.append("..")

import re
import random
from CoreSim.MScheduler import *
from CoreSim.Flow import *
from CoreSim.Event import *
from CoreSim.Link import *
from CoreSim.Unit import *


def assignNewFlow(srcHost, dstHost, flows, flowSize, startTime, mcastFlowId, numHosts, numToRs, numCores):
    flow = Flow()
    flow.startId = srcHost
    flow.endId = dstHost
    flow.SetFlowSize(flowSize)
    flow.startTime = startTime
    flow.flowId = len(flows) + 1
    flow.mcastFlowId = mcastFlowId
    flows.append(flow)
    return flow


class MSchedulerBinomial(MScheduler):
    """
    Multicast scheduler
    """

    def __init__(self, topo_args, flow_args, conf_args):
        MScheduler.__init__(self, topo_args, flow_args, conf_args)
        self.mcastSubtree = {}  # binomial tree
        self.type = 'binomial'

    def InitialMulticast(self, eventQueue):
        for mflowId in self.mcastFlows:
            mcastFlow = self.mcastFlows[mflowId]

            index = mcastFlow.groupSize / 2
            srcHost = mcastFlow.srcHost
            dstHost = mcastFlow.dstHosts.keys()[index]

            self.mcastSubtree[mflowId] = {}
            self.mcastSubtree[mflowId][srcHost] = mcastFlow.dstHosts.keys()[:index]
            self.mcastSubtree[mflowId][dstHost] = mcastFlow.dstHosts.keys()[index + 1:]
            # print srcHost, dstHost
            # print mcastFlow.dstHosts.keys()
            # print 'mcastSubtree', self.mcastSubtree

            flow = assignNewFlow(srcHost, dstHost, self.flows, mcastFlow.flowSize / 2, mcastFlow.startTime, \
                                 mcastFlow.flowId, self.numHosts, self.numToRs, self.numCores)

            # eventqueue
            event = FlowArrivalEvent(mcastFlow.startTime, flow)
            eventQueue.push(event, mcastFlow.startTime)

    def UpdateMulticast(self, mflowId, curFlow, curTime, eventQueue):
        mcastFlow = self.mcastFlows[mflowId]
        mcastFlow.dstHosts[curFlow.endId] = True

        if all(mcastFlow.dstHosts.values()):
            mcastFlow.finishTime = curTime
            for flow in self.flows:
                if flow.mcastFlowId == mflowId:
                    mcastFlow.energy += flow.energy
            print("mcastFlow:%s transmission time %s energy %s" %
                  (mflowId, mcastFlow.finishTime - mcastFlow.startTime, mcastFlow.energy))
        else:
            srcHost = curFlow.startId
            dstHost = curFlow.endId
            srcSubHosts = self.mcastSubtree[mflowId][srcHost]
            dstSubHosts = self.mcastSubtree[mflowId][dstHost]

            if srcSubHosts:
                index1 = len(srcSubHosts) / 2
                host1 = srcSubHosts[index1]
                self.mcastSubtree[mflowId][srcHost] = srcSubHosts[:index1]
                self.mcastSubtree[mflowId][host1] = srcSubHosts[index1 + 1:]

                flow = assignNewFlow(srcHost, host1, self.flows, mcastFlow.flowSize, curTime, \
                                     mcastFlow.flowId, self.numHosts, self.numToRs, self.numCores)

                event = FlowArrivalEvent(curTime, flow)
                eventQueue.push(event, event.time)

            if dstSubHosts:
                index2 = len(dstSubHosts) / 2
                host2 = dstSubHosts[index2]
                self.mcastSubtree[mflowId][dstHost] = dstSubHosts[:index2]
                self.mcastSubtree[mflowId][host2] = dstSubHosts[index2 + 1:]
                flow = assignNewFlow(dstHost, host2, self.flows, mcastFlow.flowSize, curTime, \
                                     mcastFlow.flowId, self.numHosts, self.numToRs, self.numCores)

                event = FlowArrivalEvent(curTime, flow)
                eventQueue.push(event, event.time)
